/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: 集合通信算子信息库头文件
 * Author: jiangchen
 * Create: 2020-5-15
 */

#ifndef OPS_KERNEL_BUILDER__BASEH
#define OPS_KERNEL_BUILDER__BASEH

#include "comm.h"
#include "base.h"
#include "common/opskernel/ge_task_info.h"
#include "common/opskernel/ops_kernel_info_types.h"
#include "common/opskernel/ops_kernel_builder.h"
#include "hcom_ops_stores.h"
#include "graph/node.h"
#include "proto/task.pb.h"
#include "op_common_def.h"


namespace hccl {
// Ge适配的类
class HCCLOpsKernelBuilder : public ge::OpsKernelBuilder {
public:
    HCCLOpsKernelBuilder();
    ~HCCLOpsKernelBuilder() override;
    // initialize opsKernelInfoStore
    ge::Status Initialize(const map<string, string> &options) override;
    // close opsKernelInfoStore
    ge::Status Finalize() override;
    // memory allocation requirement
    virtual ge::Status CalcOpRunningParam(ge::Node& node) = 0;
    // generate taskinfo for op
    virtual ge::Status GenerateTask(const ge::Node &node, ge::RunContext &runContext,
        std::vector<domi::TaskDef> &taskDefList) = 0;
protected:
    virtual HcclResult GetSupportedOP(std::vector<std::string> &hcclSupportOp) const
    {
        return HCCL_SUCCESS;
    };
    virtual HcclResult SetOpMemAttr(ge::Node& node, const std::string &sCollectiveType, const u64 &opMemSize)
    {
        return HCCL_SUCCESS;
    };
    virtual HcclResult SetOpAtomicInputIndex(ge::Node& node, const std::string &sCollectiveType)
    {
        return HCCL_SUCCESS;
    };
    HcclResult CheckSupportedOP(const std::string &sCollectiveType) const;
    HcclResult SetOpOutputMemSize(ge::Node& node, const std::string &sCollectiveType);
    HcclResult CalcHCCLOutputMemSize(const std::string &sCollectiveType, int64_t& memSize);
};
}
#endif // __OPS_KERNEL_INFO_STORE__BASEH__
